import pandas as pd
import pyreadstat
from dataclasses import dataclass
from functools import cached_property
from typing import Dict

@dataclass(kw_only=True)
class CharlsEncoding:
    new_code: int
    meaning: str

class CharlsProcessor:
    def __init__(self, dta_file_path: str = None, codebook_path: str = None):
        """
        Initializes the CHARLS processor.

        Args:
            dta_file_path (str, optional): Path to the .dta file. Defaults to None.
            codebook_path (str, optional): Path to the codebook.csv file. Defaults to None.
        """
        self.dta_file_path = dta_file_path
        self.codebook_path = codebook_path
        if self.codebook_path:
            self.codebook = pd.read_csv(self.codebook_path)

    @cached_property
    def encodings(self) -> Dict[int, CharlsEncoding]:
        """
        Creates a mapping from new_code to the actual categorical content.

        The mapping data is obtained from the codebook.csv file.
        Handles possible encoding exceptions to ensure the accuracy of the mapping relationship.

        Returns:
            Dict[int, CharlsEncoding]: A dictionary mapping new_code to CharlsEncoding objects.
        """
        encodings_map = {}
        if self.codebook is not None:
            for _, row in self.codebook.iterrows():
                try:
                    new_code = int(row['new_code'])
                    meaning = str(row['label'])
                    encodings_map[new_code] = CharlsEncoding(new_code=new_code, meaning=meaning)
                except (ValueError, KeyError) as e:
                    print(f"Skipping row due to error: {e} - Row: {row}")
        return encodings_map

    @cached_property
    def columns_name_to_description(self) -> Dict[str, str]:
        """
        Gets the mapping relationship between column names and labels.

        The mapping data is extracted from the codebook of the .dta file.
        Ensures that the .dta file format is handled correctly to extract accurate column name description information.

        Returns:
            Dict[str, str]: A dictionary mapping column names to their descriptions.
        """
        columns_name_to_description = {}
        if self.dta_file_path:
            try:
                _, meta = pyreadstat.read_dta(self.dta_file_path, metadataonly=True)
                for var_name, var_label in meta.column_names_to_labels.items():
                    columns_name_to_description[var_name] = var_label
            except Exception as e:
                print(f"An error occurred while reading the .dta file: {e}")
        return columns_name_to_description

    def get_column_data_types(self) -> Dict[str, str]:
        """
        Gets the mapping between column names and their data types.

        For each variable, it first checks if a codebook classification exists (i.e., has value labels).
        If it does, the variable's data type is marked as "category".
        Otherwise, the data type is extracted from readstat_variable_types.

        The mapping data is extracted from the metadata of the .dta file.

        Returns:
            Dict[str, str]: A dictionary mapping column names to their data types.
        """
        if not self.dta_file_path:
            return {}
        try:
            _, meta = pyreadstat.read_dta(self.dta_file_path, metadataonly=True)
            data_types = meta.readstat_variable_types
            # If a variable has value labels, it is considered categorical.
            if hasattr(meta, 'variable_value_labels'):
                for var_name in meta.variable_value_labels:
                    if var_name in data_types:
                        data_types[var_name] = "category"
            return data_types
        except Exception as e:
            print(f"An error occurred while reading the .dta file: {e}")
            return {}